package com.ams.server.replication;

import java.io.IOException;
import java.util.HashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.ams.protocol.rtmp.RtmpConnection;
import com.ams.protocol.rtmp.RtmpException;
import com.ams.protocol.rtmp.RtmpHeader;
import com.ams.protocol.rtmp.amf.AmfException;
import com.ams.protocol.rtmp.amf.AmfValue;
import com.ams.protocol.rtmp.message.RtmpMessage;
import com.ams.protocol.rtmp.message.RtmpMessageCommand;
import com.ams.protocol.rtmp.net.NetConnectionException;
import com.ams.protocol.rtmp.net.StreamPublisher;
import com.ams.protocol.rtmp.net.StreamPublisherManager;

public class ReplSubscriber {
    final private Logger logger = LoggerFactory.getLogger(ReplSubscriber.class);

    private RtmpConnection rtmp;
    private HashMap<Integer, String> streams = new HashMap<Integer, String>();

    public ReplSubscriber(RtmpConnection rtmp) {
        this.rtmp = rtmp;
    }

    public void sendSubscriptionCommand(String streamName) throws IOException {
        AmfValue[] args = AmfValue.array(streamName);
        RtmpMessage message = new RtmpMessageCommand("subscribe", 0, args);
        rtmp.writeRtmpMessage(0, 0, 0, message);
    }

    public void acceptPublisherCommand() throws NetConnectionException,
            IOException, RtmpException, AmfException {
        // publisher will send publish command to subscriber
        if (!rtmp.readRtmpMessage())
            return;
        RtmpHeader header = rtmp.getCurrentHeader();
        RtmpMessage message = rtmp.getCurrentMessage();

        switch (message.getType()) {
        case RtmpMessage.MESSAGE_AMF0_COMMAND:
            RtmpMessageCommand command = (RtmpMessageCommand) message;
            if ("publish".equals(command.getName())) {
                int streamId = header.getStreamId();
                AmfValue[] args = command.getArgs();
                String publishName = args[1].string();
                streams.put(streamId, publishName);
                if (StreamPublisherManager.getPublisher(publishName) == null) {
                    StreamPublisherManager
                            .addPublisher(new StreamPublisher(publishName));
                }
                logger.info("received publish command: {}", publishName);
            } else if ("closeStream".equals(command.getName())) {
                int streamId = header.getStreamId();
                String publishName = streams.get(streamId);
                if (publishName == null)
                    break;
                streams.remove(streamId);
                StreamPublisher publisher = StreamPublisherManager
                        .getPublisher(publishName);
                if (publisher != null) {
                    publisher.close();
                    StreamPublisherManager.removePublisher(publishName);
                    logger.info("received close stream command: {}",
                            publishName);
                }
            }

            break;
        case RtmpMessage.MESSAGE_AUDIO:
        case RtmpMessage.MESSAGE_VIDEO:
        case RtmpMessage.MESSAGE_AMF0_DATA:
            int streamId = header.getStreamId();
            String publishName = streams.get(streamId);
            if (publishName == null)
                break;
            // publish to local subscriber
            StreamPublisher publisher = StreamPublisherManager
                    .getPublisher(publishName);
            if (publisher != null) {
                publisher.publish(message.toMediaMessage(header.getTimestamp()));
            }
            break;
        case RtmpMessage.MESSAGE_ACK:
            logger.debug("received publisher ping");
            break;
        }
    }
}